Add support for async cuda copies from global to shared - #9263
Conversation
A store into shared memory whose value is a plain load from global memory is emitted as an asynchronous copy on sm_80 and later, which moves the data without routing it through registers. The copies issued in a producer are waited for at the end of it. This takes the shared-memory tensor core matmul from 31.4 to 39.5 TFlop/s at 2048^3 on an RTX 5060 Ti. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
If nothing between a barrier and the next one at the same level loads what was stored before it, the later barrier can do the job of both. The fence types of the elided barrier are added to the one that remains. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
A Func stored in GPUSharedAsync goes in GPU shared memory, but is written by an asynchronous copy instruction that moves the data straight from global memory without routing it through registers. Asking for the memory type is a promise that every store to it is a copy the hardware can make that way, so a store that isn't reports what the copy engine requires and how a schedule usually satisfies it, rather than quietly falling back to a load and a store. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
AMXTile was never added, and GPUSharedAsync is new. List them in the same order as the enum in Expr.h so that a missing one is easier to spot. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Cover each of the three copy widths the hardware supports at several element sizes, plus the shapes a staged input takes: a two-dimensional tile, a padded stride, two inputs staged into one kernel, and a Func::in wrapper. The error test checks that breaking each constraint produces a user error saying which one, rather than crashing or quietly falling back to a load and a store. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
The peephole fired on any store into shared memory that matched, so a Func in plain GPUShared got an asynchronous copy too, and there was no way to ask for the synchronous version. That made the memory type only control whether a store that didn't match was an error, rather than whether the copy was asynchronous at all. Require the destination to be GPUSharedAsync, so that GPUShared and GPUSharedAsync are the two ways to ask for the two lowerings, and a schedule can compare them. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
A store whose value wasn't a plain load was always reported as the Func not being a copy, but that also caught two cases where it is a copy and something else is wrong. A source read with a stride is broken into a shuffle of dense loads before it reaches here, and a source computed elsewhere in the kernel is a load from the wrong place. Report those as themselves. Also test the predicated case, which needs an align_storage to keep the destination aligned, or the alignment requirement is what fails first. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
The density of the two ends of the copy is checked by strided_ramp_base, whose default stride of one is doing the work. That is easy to misread as extracting an address, so say so where it is called. A source read with a stride never reaches that check, because it is broken into a shuffle of dense loads first and fails the earlier test for being a plain load. Storing the staged Func in the opposite order to the one it is read in is what reaches it, so test that. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
|
What bugs me a little bit with this approach is that it's not actually a different type of memory. It's a different copy-operation. It reminds me more of the streaming-store and streaming-loads we did two weeks ago. I wonder if it wouldn't make more sense to do something like: func.compute_at(other, block).gpu_thread(_0, _1).streaming_load().streaming_store();I realize they are not "streaming" or non-temporal in the sense that the CPU-instructions are non-temporal skipping cache-writebacks/loads, but at least the scheduling language is clear that we are modifying the store/load operations instead of the memory type. |
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## main #9263 +/- ##
==========================================
+ Coverage 69.79% 69.89% +0.10%
==========================================
Files 258 258
Lines 77956 78206 +250
Branches 18979 19036 +57
==========================================
+ Hits 54410 54665 +255
- Misses 17803 17818 +15
+ Partials 5743 5723 -20 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
A kernel could only ever have one shared memory allocation, because the PTX backend gives them all base address zero in address space 3. Adding GPUSharedAsync broke that: shared allocations are clustered by memory type, so a kernel that staged one Func synchronously and another asynchronously got two clusters, and they overlapped. The consumer read whatever the other copy had written. Being filled by the copy engine is a property of the stores to an allocation, not of the memory it lives in, so lower it as one. A pass inside fuse_gpu_thread_loops rewrites GPUSharedAsync allocations to ordinary shared memory, wraps the value of each store to them in a cuda_bypass_registers intrinsic, and awaits the copies at the end of the producer. The allocations then all cluster together again, which also lets them share space by lifetime the way they always could within a memory type. The intrinsic carries the group its copies belong to, so that consuming one Func doesn't wait for the copies into another. Waits are FIFO - cp.async can only wait for all but the newest N groups - so the backend tracks the committed groups and lowers a group to its position in that list. A barrier still waits for everything, because it publishes shared memory to the rest of the block. Since allocations are now packed together whatever their memory type, a group has to end on a 16 byte boundary for whatever follows it to be alignable enough to copy into, and the intrinsic carries the Func name so an error can name the Func rather than the packed allocation. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
The Makefile globs the test directory, so these have been running there since they were written, but CMake lists each test explicitly and they were never added. That is what CI builds from. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Eliding a barrier by widening a later one is an optimization for kernels with more than one shared allocation, and has nothing to do with asynchronous copies - it fires just as often on kernels that have none. It also arrived with no tests, and a barrier wrongly elided is a race rather than a wrong answer, so it needs its own. The wait for outstanding copies at a barrier stays here, because a barrier publishing shared memory to the rest of the block does require that any copies into it have landed. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
|
@mcourteaux — I think we ought to model it as a different type of memory. As formal functions, Load : forall M. T*[M] -> T[Reg]
Store : forall M. T[Reg] -> T*[M] -> unitWhere a type For a simple, existing, example of the distinction I'm trying to draw, consider the prefetch: Prefetch : forall M. T*[M] -> unit
We have a similar situation here with CopyAsync : T*[Global] -> T*[Shared] -> T*[SharedAsync]
WaitGroup : T*[SharedAsync] -> T*[Shared]You can sort of think of the return value of |
You make a compelling argument. It's a matter of modeling choices indeed. Was just a thought of mine. Seems like Andrew already instructed Claude to remodel based on my comment 😬. |
|
I was actually just fixing a lowering issue, and didn't change the front-end syntax, so no work was wasted. |
I wonder if using ring_buffer() on these copies will work here and make them asynchronous (when I added it I sort of looked at a similar case, but with DMA on DSP). |
|
Yes! My plan is to make ring_buffer() work on them, and to make async() map to warp specialization. I have another branch that does this. However I also want to be able to do it non-async - just a fixed depth of software pipelining. I think software pipelining would be sliding window with the option to compute the producer a few iterations ahead. |
align_up is now instantiated with Expr as well as integral types, and clang-tidy's performance-unnecessary-value-param fires on the by-value parameter for types with non-trivial copies. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
# Conflicts: # src/FuseGPUThreadLoops.cpp
| // an expression that stands alone, which is what the analysis below and | ||
| // codegen_buffer_pointer both need. It has to be held in a local, because | ||
| // everything below points into it. | ||
| const Expr stored = substitute_in_all_lets(op->value); |
There was a problem hiding this comment.
This should probably peel lets, not substitute them. Consider a help to peel lets because we do it in lots of places.
| return false; | ||
| } | ||
| } | ||
|
|
There was a problem hiding this comment.
lets will need to be rewrapped, or at least placed in scope, here
|
|
||
| { | ||
| const char *reason = ""; | ||
| if (codegen_async_copy(op, &reason)) { |
There was a problem hiding this comment.
This looks like it operates unconditionally, instead of just for allocations marked GPUSharedAsync
| // Asking for that memory type is a promise that the stores to it are | ||
| // copies the hardware can make asynchronously. If one isn't, say so | ||
| // rather than quietly emitting a load and a store instead. | ||
| const Expr stored = substitute_in_all_lets(op->value); |
There was a problem hiding this comment.
It's weird to have to re-check for the cuda_bypass_registers call just to see if it failed. Why not just check the memory type?
| if (unit_bytes >= alignment) { | ||
| return size; | ||
| } | ||
| return align_up(size, alignment / unit_bytes); |
There was a problem hiding this comment.
Consider just changing align_up to return the first arg if the second arg is zero. Also, please assert unit_bytes is a power of two somewhere.
|
|
||
| // Run `body` and assert it produces a Halide user error mentioning `substring`. | ||
| template<typename F> | ||
| bool expect_user_error(const char *name, const char *substring, F body) { |
There was a problem hiding this comment.
Dedup with other error tests
alexreinking
left a comment
alexreinking
left a comment
There was a problem hiding this comment.
Review done in-person with @abadams
The wait was appended after the producer, which put it outside the loops over threads, and a statement outside those loops is one that only the first thread runs. Waiting for a copy is something each thread does for the copies it issued itself, so the wait belongs at the end of the innermost loop over threads. The generated code was right anyway, because the backend emits the wait from its own bookkeeping rather than from the intrinsic, so the guard never reached the PTX. That is a thin thing to rely on: were the backend to honour it, every thread but the first would read the staged data without waiting for its own copies to land. Putting it in the right place also lets the loops over threads of two staged inputs fuse, so both sets of copies are now in flight before either is waited on. apps/cuda_mat_mul goes from 172us to 160us against cublas's 150. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> (cherry picked from commit e62a4f4f259415abc906489062c37390682ced13)
CSE and LICM lift common subexpressions of a stored value into lets around it, which hides the marker saying the store should be an asynchronous copy. Substituting them back in got past that, but it repeated each value at every use, which is what lifting them out avoided, so the addresses were built from expressions that had lost their sharing. Peel the lets instead, and put them in scope while the addresses are built. Peeling lets is done in a lot of places, so add a helper for it. Detecting the marker also decided two separate things: whether this store was one to emit as a copy, and, once that had failed, whether the schedule had asked for a copy at all and so should be told why it didn't get one. The second reading meant looking for the marker again in the caller. Report the two separately instead, and test for the marker first, so that an ordinary store costs a look at the value and nothing more. The allocation's memory type can't answer either question here: MemoryType::GPUSharedAsync is rewritten to ordinary shared memory long before codegen, which is what lets shared allocations be packed together. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Group sizes are counted in units of the type they are allocated as, so the number of units to align to is 16 bytes divided by the size of one. Units of 16 bytes or more need no alignment and give zero, which was a special case in a helper. Have align_up treat an alignment of zero as no alignment at all, and the helper is a single expression, so drop it. Assert that allocation types are a power of two bytes wide, next to the existing check that they divide each other, which covers both of the types these sizes are converted between. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Two tests had near-identical copies. Keep the one that checks the message mentions what the test broke, so that a test can't pass on some unrelated error, and give it an overload for callers that don't care what the message says. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
# Conflicts: # test/correctness/tiled_matmul_errors.cpp
4 participants
CUDA has instructions that copy data directly from global into shared. This PR adds the ability to ask for them, via a new MemoryType::GPUSharedAsync, which means store in shared, and write only using async write instructions. They're not especially asynchronous right now (there's a barrier almost immediately), because we have no software pipelining, but they at least don't consume valuable registers just to move memory around. I have confirmed this is a small performance win on a split-k schedule for cuda mat mul, but I haven't committed that because it's worse than the current schedule. This really matters for a future tensorcore PR, where staging in shared is more important.
Drive-by fix to add AMXTile to PyEnums.cpp